#!/usr/bin/env bash

###############################################################################
# Formats ipython notebooks with tensorflow-docs nbformat tool.
#
# Usage:
#     check/nbformat [--apply]
#
# Without '--apply', the diff that would be applied is printed and the exit
# status is 1 if there are any changes or else 0 if no changes are needed.
#
# With '--apply', the exit status is 0 and the changed files are actually
# reformatted.
#
################################################################################

only_print=1
for arg in "$@"; do
    if [[ "${arg}" == "--apply" ]]; then
        only_print=0
    else
        echo -e "\033[31mToo many arguments. Expected [--apply].\033[0m" >&2
        exit 1
    fi
done

# Get the working directory to the repo root.
thisdir=$(dirname "${BASH_SOURCE[0]:?}") || exit $?
repo_dir=$(git -C "${thisdir}" rev-parse --show-toplevel) || exit $?
cd "${repo_dir}" || exit $?

# check if the nbfmt tool in tensorflow-docs is installed and working
if ! python3 -c "import tensorflow_docs.tools.nbfmt"; then
    echo -e "\033[31mInstallation of tensorflow-docs is broken, please correct.\033[0m"
    exit 2
fi

my_nbfmt() {
    python3 -m tensorflow_docs.tools.nbfmt --indent=1 "$@"
}

exit_code=0

# Check if notebooks are already formatted and store output for later use
if output=$(my_nbfmt --test docs 2>&1); then
    echo -e "\033[32mNotebooks are formatted\033[0m."
elif (( only_print )); then
    echo -e "\033[31mThe following notebooks require formatting\033[0m."
    grep "^- docs" <<<"${output}"
    echo -e "\033[33mNotebooks are not formatted. Please run 'check/nbformat --apply'\033[0m"
    exit_code=1
else
    my_nbfmt docs &&
        echo -e "\033[33mReformatted changed notebooks\033[0m." ||
        exit_code=$?
fi

exit ${exit_code}
